#include "nnslitemtree.h"
#include "nnslitememory.h"

static nnsliteMTreeNode* NodeCreate(nnsliteMTree* tree, bool leaf);
static NNSLITE_RET NodeFree(nnsliteMTreeNode* node);
static NNSLITE_RET NodeZero(nnsliteMTreeNode* node);
static NNSLITE_RET NodeWrite(nnsliteMTreeNode* node, nnsliteBufferedHeapFile* bhf);
static NNSLITE_RET NodeRead(nnsliteMTreeNode* node, nnsliteBufferedHeapFile* bhf);
static NNSLITE_RET FindSplitSeeds(nnsliteMTree* tree, u32 nodeindex, nnsliteObject* seed1, nnsliteObject* seed2);
static NNSLITE_RET NodeAbsorb(nnsliteMTree* tree, u32 nodeindex, nnsliteMTreeNode* seedNode1, nnsliteMTreeNode* seedNode2);
static NNSLITE_RET NodeSplit(nnsliteMTree* tree, u32 nodeindex);
static NNSLITE_RET NodeInsert(nnsliteMTree* tree, u32 nodeindex, const nnsliteObject* object);
static NNSLITE_RET NodeNNSearch(nnsliteMTree* tree, u32 nodeindex, const nnsliteObject* object, u32 top, nnsliteObject* results, f64* resultDists);
static NNSLITE_RET NodePrint(nnsliteMTreeNode* node, FILE* fp, u32 indent);
static NNSLITE_RET NodeRemove(nnsliteMTreeNode* node);

/** nnsliteMTreeNode  */
static NNSLITE_RET NodeZero(nnsliteMTreeNode* node)
{
	memset(node, 0, sizeof(nnsliteMTreeNode));
	return NNSLITE_RET_OK;
}

static nnsliteMTreeNode* NodeCreate(nnsliteMTree* tree, bool leaf)
{
	NNSLITE_RET ret;
	u32 i;
	nnsliteMTreeNode* node;
	/* Find a unused node */
	for (i=0; i<tree->maxNumNodes; i++)
	{
		if (tree->nodes[i].isUsed == false)
			break;
	}
	if (i == tree->maxNumNodes)
	{
		/* No available node */
		assert(false);
		return NULL;
	}
	node = &tree->nodes[i];
	node->isLeaf = leaf;
	node->isUsed = true;
	node->pTree = tree;
	node->nodeIndex = i;
	node->parentIndex = U32_NULL;
	node->childCount = 0;
	node->depth = 0;
	node->radius = 0.0;
	nnsliteObjectInit(&node->object);
	if (leaf)
	{
		/* Create a new object container for this leaf node*/
		ret = nnsliteObjectContainerCreate(tree->objContainerBase, &node->objContainerID);
		if (NNSLITE_FAILED(ret))
		{
			assert(false);
			return NULL;
		}
	}
	else
	{
		node->objContainerID = U32_NULL;
	}
	return node;
}

static NNSLITE_RET NodeFree(nnsliteMTreeNode* node)
{
	NNSLITE_RET ret;
	nnsliteMTree* tree = node->pTree;
	nnsliteObjectContainer* container;
	
	nnsliteObjectFree(&node->object);
	if (node->isLeaf)
	{
		container = nnsliteObjectContainerGet(tree->objContainerBase, node->objContainerID);
		/* Close the object container of this leaf node */
		ret = nnsliteObjectContainerClose(container);
		NNSLITE_ON_FAILED_RETURN(ret);
	}
	return NNSLITE_RET_OK;
}

static NNSLITE_RET NodeRemove(nnsliteMTreeNode* node)
{
	NNSLITE_RET ret;
	nnsliteMTree* tree = node->pTree;
	
	if (node->isLeaf)
	{
		/* Delete the object container of this leaf node*/
		ret = nnsliteObjectContainerDelete(tree->objContainerBase, node->objContainerID);
		NNSLITE_ON_FAILED_RETURN(ret);
	}
	/* Free this node first */
	ret = NodeFree(node);
	NNSLITE_ON_FAILED_RETURN(ret);
	/* Zero this node */
	NodeZero(node);
	return NNSLITE_RET_OK;
}

static NNSLITE_RET NodeWrite(nnsliteMTreeNode* node, nnsliteBufferedHeapFile* bhf)
{
	nnsliteBufferedHeapFileWrite(bhf, node, sizeof(nnsliteMTreeNode));
	if (node->object.size > 0)
		nnsliteBufferedHeapFileWrite(bhf, node->object.data, node->object.size);
	return NNSLITE_RET_OK;
}

static NNSLITE_RET NodeRead(nnsliteMTreeNode* node, nnsliteBufferedHeapFile *bhf)
{
	nnsliteBufferedHeapFileRead(bhf, node, sizeof(nnsliteMTreeNode));
	if (node->object.size > 0)
	{
		node->object.data = (u8*)nnsliteMAlloc(node->object.size);
		nnsliteBufferedHeapFileRead(bhf, node->object.data, node->object.size);
	}
	else
	{
		node->object.data = NULL;
	}
	return NNSLITE_RET_OK;
}

static NNSLITE_RET FindSplitSeeds(nnsliteMTree* tree, u32 nodeindex, nnsliteObject* seed1, nnsliteObject* seed2)
{
	s32 ret;
	f64 maxDist = 0.0;
	f64 dist;
	nnsliteMTreeNode* node = &tree->nodes[nodeindex];
	if (node->isLeaf)
	{
		u32 i;
		nnsliteObject obj;
		nnsliteObjectContainerBase* objContainerBase;
		nnsliteObjectContainer* container;
		
		objContainerBase = tree->objContainerBase;
		container = nnsliteObjectContainerGet(objContainerBase, node->objContainerID);
		/* Begin to scan each object of the container */
		ret = nnsliteObjectContainerBeginScan(container);
		NNSLITE_ON_FAILED_RETURN(ret);
		/* Initialize the temporary object */
		nnsliteObjectInit(&obj);
		for (i=0;i<container->objCount; i++)
		{
			/* Get the next object in the container */
			ret = nnsliteObjectContainerNext(container, &obj);
			NNSLITE_ON_FAILED_RETURN(ret);
			/* Calculate the distance */
			dist = tree->distfunc(&node->object, &obj);
			if (dist > maxDist)
			{
				nnsliteObjectCopy2(seed2, &obj);
			}
		}
		*seed1 = nnsliteObjectCopy(&node->object);

		nnsliteObjectFree(&obj);
	}
	else
	{
		nnsliteMTreeNode* child;
		u32 childindex;
		u32 i;
		for(i=0;i<node->childCount;i++)
		{
			childindex = node->childIndexes[i];
			child = &tree->nodes[childindex];
			dist = tree->distfunc(&node->object, &child->object);
			if (dist> maxDist)
			{
				nnsliteObjectCopy2(seed2, &child->object);
			}
		}
		*seed1 = nnsliteObjectCopy(&node->object);
	}

	return NNSLITE_RET_OK;
}

static NNSLITE_RET NodeAbsorb(nnsliteMTree* tree, u32 nodeindex, nnsliteMTreeNode* seedNode1, nnsliteMTreeNode* seedNode2)
{
	s32 ret;
	nnsliteMTreeNode* parentNode;
	nnsliteObject obj;
	nnsliteObject* seed1 = &seedNode1->object;
	nnsliteObject* seed2 = &seedNode2->object;
	f64 dist1,dist2;

	parentNode = &tree->nodes[nodeindex];
	if (parentNode->isLeaf)
	{
		u32 i;
		nnsliteObjectContainerBase* objContainerBase;
		nnsliteObjectContainer *parentContainer,*container1,*container2;

		/* Get the containers of the 3 nodes */
		objContainerBase = tree->objContainerBase;
		parentContainer = nnsliteObjectContainerGet(objContainerBase, parentNode->objContainerID);
		container1 = nnsliteObjectContainerGet(objContainerBase, seedNode1->objContainerID);
		container2 = nnsliteObjectContainerGet(objContainerBase, seedNode2->objContainerID);
		/* Begin to scan each object of the container */
		ret = nnsliteObjectContainerBeginScan(parentContainer);
		NNSLITE_ON_FAILED_RETURN(ret);
		/* Initialize the object */
		ret = nnsliteObjectInit(&obj);
		NNSLITE_ON_FAILED_RETURN(ret);

		for (i=0;i<parentContainer->objCount;i++)
		{
			/* Get the next object in the container */
			ret = nnsliteObjectContainerNext(parentContainer, &obj);
			NNSLITE_ON_FAILED_RETURN(ret);

			dist1 = tree->distfunc(&obj, seed1);
			dist2 = tree->distfunc(&obj, seed2);
			if (dist1 < dist2)
			{
				nnsliteObjectContainerInsert(container1, &obj);
				seedNode1->radius = MAX(seedNode1->radius, dist1);
			}
			else
			{
				nnsliteObjectContainerInsert(container2, &obj);
				seedNode2->radius = MAX(seedNode2->radius, dist2);
			}
		}
		nnsliteObjectFree(&obj);
	}
	else
	{
		u32 i;
		seedNode1->radius = 0.0;
		seedNode2->radius = 0.0;
		for (i=0;i<parentNode->childCount; i++)
		{
			nnsliteMTreeNode* childNode;
			u32 childIndex = parentNode->childIndexes[i];
			childNode = &tree->nodes[childIndex];
			dist1 = tree->distfunc(&childNode->object, seed1);
			dist2 = tree->distfunc(&childNode->object, seed2);
			if (dist1 < dist2)
			{
				seedNode1->childIndexes[seedNode1->childCount++] = childIndex;
				seedNode1->radius = MAX(seedNode1->radius, dist1+childNode->radius);
				childNode->parentIndex = seedNode1->nodeIndex;
			}
			else
			{
				seedNode2->childIndexes[seedNode2->childCount++] = childIndex;
				seedNode2->radius = MAX(seedNode2->radius, dist2+childNode->radius);
				childNode->parentIndex = seedNode2->nodeIndex;
			}
		}
	}

	return NNSLITE_RET_OK;
}

static NNSLITE_RET NodeSplit(nnsliteMTree* tree, u32 nodeindex)
{
	NNSLITE_RET ret;
	nnsliteObject seed1, seed2;
	nnsliteMTreeNode* seedNode1;
	nnsliteMTreeNode* seedNode2;
	nnsliteMTreeNode* node;
	nnsliteMTreeNode* parentNode;

	/* get the seeds to split */
	nnsliteObjectInit(&seed1);
	nnsliteObjectInit(&seed2);
	ret = FindSplitSeeds(tree, nodeindex, &seed1, &seed2);
	NNSLITE_ON_FAILED_RETURN(ret);

	/* get the node*/
	node = &tree->nodes[nodeindex];
	
	/** create two new nodes */
	seedNode1 = NodeCreate(tree, node->isLeaf);
	seedNode2 = NodeCreate(tree, node->isLeaf);
	seedNode1->object = nnsliteObjectCopy(&seed1);
	seedNode2->object = nnsliteObjectCopy(&seed2);
	seedNode1->depth = node->depth;
	seedNode2->depth = node->depth;
	seedNode1->isLeaf = node->isLeaf;
	seedNode2->isLeaf = node->isLeaf;

	/* absorbs the nodes by the seeds */
	ret = NodeAbsorb(tree, nodeindex, seedNode1, seedNode2);
	NNSLITE_ON_FAILED_RETURN(ret);

	if (node->parentIndex != U32_NULL) /* has parent node */
	{
		u32 i;
		parentNode = &tree->nodes[node->parentIndex];
		/* Replace the original node indexx by the seed node 1's */
		for (i=0;i<parentNode->childCount;i++)
		{
			if (parentNode->childIndexes[i] == node->nodeIndex)
			{
				parentNode->childIndexes[i] = seedNode1->nodeIndex;
				break;
			}
		}
		/* Add the seed node 2 to the child indexes */
		parentNode->childIndexes[parentNode->childCount++] = seedNode2->nodeIndex;
		seedNode1->parentIndex = node->parentIndex;
		seedNode2->parentIndex = node->parentIndex;
		/* add the two seed nodes */
		NodeRemove(node);
		/* check the condition for the splitting of parent node */
		if (parentNode->childCount == tree->maxChildCount)
		{
			ret = NodeSplit(tree, parentNode->nodeIndex);
			NNSLITE_ON_FAILED_RETURN(ret);
		}
	}
	else /* does not have parent node */
	{
		u32 i;
		nnsliteMTreeNode* newRootNode;
		/* update the depth of all nodes */
		for(i=0;i<tree->maxNumNodes;i++)
		{
			if (tree->nodes[i].isUsed)
			{
				tree->nodes[i].depth++;
			}
		}
		seedNode1->depth++;
		seedNode2->depth++;
		/* initialize the new root node, which mustn's be a leaf node.
		 * So there's no need to create object cotainer for it */
		newRootNode = NodeCreate(tree, false);
		newRootNode->depth = 0;
		newRootNode->isLeaf = FALSE;
		newRootNode->childIndexes[newRootNode->childCount++] = seedNode1->nodeIndex;
		newRootNode->childIndexes[newRootNode->childCount++] = seedNode2->nodeIndex;
		newRootNode->object = nnsliteObjectCopy(&node->object);
		newRootNode->radius = node->radius;
		/* add the two seed nodes */
		seedNode1->parentIndex = newRootNode->nodeIndex;
		seedNode2->parentIndex = newRootNode->nodeIndex;
		NodeRemove(node);
		/* update the new root node index of the tree */
		tree->rootNodeIndex = newRootNode->nodeIndex;
	}
	
	nnsliteObjectFree(&seed1);
	nnsliteObjectFree(&seed2);
	
	return NNSLITE_RET_OK;
}

static NNSLITE_RET NodeInsert(nnsliteMTree* tree, u32 nodeindex, const nnsliteObject* object)
{
	s32 ret;
	f64 distToRoute;
	nnsliteMTreeNode* node;
	nnsliteDistFunc distfunc = tree->distfunc;

	/* get the node*/
	node = &tree->nodes[nodeindex];
	if (NNSLITE_OBJ_ISNULL(&node->object))
	{
		node->object = nnsliteObjectCopy(object);
		node->radius = 0.0;
	}
	
	distToRoute = distfunc(&node->object, object);
	node->radius = MAX(node->radius, distToRoute);
	
	if (node->isLeaf)
	{
		nnsliteObjectContainerBase* objContainerBase = tree->objContainerBase;
		nnsliteObjectContainer* container = nnsliteObjectContainerGet(objContainerBase, node->objContainerID);
		/* Insert the object to the node's container */
		nnsliteObjectContainerInsert(container, object);
		if (container->objCount >= tree->maxChildCount && node->depth < tree->maxDepth)
		{
			ret = NodeSplit(tree, nodeindex);
			NNSLITE_ON_FAILED_RETURN(ret);
		}
	}
	else
	{
		f64 d;
		f64 minDist = (f64)DBL_MAX;
		u32 i;
		nnsliteMTreeNode* childNode;
		nnsliteMTreeNode* nearestChildNode = NULL;
		for (i=0; i<node->childCount; i++)
		{
			u32 childIndex = node->childIndexes[i];
			childNode = &tree->nodes[childIndex];
			d = distfunc(&childNode->object, object);
			if (d < minDist)
			{
				minDist = d;
				nearestChildNode = childNode;
			}
		}

		ret = NodeInsert(tree, nearestChildNode->nodeIndex, object);
		NNSLITE_ON_FAILED_RETURN(ret);
	}
	return NNSLITE_RET_OK;
}

static NNSLITE_RET NodeNNSearch(nnsliteMTree* tree, u32 nodeindex, const nnsliteObject* object, u32 top, nnsliteObject* results,f64* resultDists)
{
	s32 ret;
	u32 i;
	f64 dist;
	nnsliteMTreeNode* node;
	/* Get the node */
	node = &tree->nodes[nodeindex];
	if (node->isLeaf)
	{
		u32 i,j,pos;
		nnsliteObject obj;
		nnsliteObjectContainerBase* objContainerBase;
		nnsliteObjectContainer* container;

		objContainerBase = tree->objContainerBase;
		container = nnsliteObjectContainerGet(objContainerBase, node->objContainerID);
		/* Begin to scan each object of the container */
		ret = nnsliteObjectContainerBeginScan(container);
		NNSLITE_ON_FAILED_RETURN(ret);
		/* Initialize the object */
		ret = nnsliteObjectInit(&obj);
		NNSLITE_ON_FAILED_RETURN(ret);
		
		for(i=0;i<container->objCount; i++)
		{
			/* Get the next object in the container */
			ret = nnsliteObjectContainerNext(container, &obj);
			NNSLITE_ON_FAILED_RETURN(ret);
			/* Calculate the distance between objects */
			dist = tree->distfunc(object, &obj);
			for (j=top-1;j>=0;j--)
			{
				if (dist > resultDists[j])
					break;
			}
			pos = (u32)(j+1);
			if (pos < top)
			{
				/* An example : 0 1 2 3 4  top=5,j=2 */
				nnsliteObjectFree(&results[top-1]); /* Tick out the last one */
				if (pos < top-1)
				{
					memmove(&results[pos+1],&results[pos], (top-pos-1)*sizeof(nnsliteObject));
					memmove(&resultDists[pos+1], &resultDists[pos], (top-pos-1)*sizeof(f64));
				}
				results[pos] = nnsliteObjectCopy(&obj);
				resultDists[pos] = dist;
			}
		}

		nnsliteObjectFree(&obj);
	}
	else
	{
		nnsliteMTreeNode* childNode;
		u32 childIndex;
		for (i=0; i<node->childCount; i++)
		{
			childIndex = node->childIndexes[i];
			childNode = &tree->nodes[childIndex];
			dist = tree->distfunc(object, &childNode->object);
			if (dist - childNode->radius < resultDists[top-1])
			{
				ret = NodeNNSearch(tree, childIndex, object, top, results, resultDists);
				NNSLITE_ON_FAILED_RETURN(ret);
			}
		}
	}

	return NNSLITE_RET_OK;
}


NNSLITE_RET nnsliteMTreeCreate(nnsliteMTree** tree, u32 maxnodes, u32 maxdepth, u32 childcount, 
		nnsliteObjectContainerBase* objContainerBase, nnsliteDistFunc func)
{
	u32 i;
	nnsliteMTreeNode* root;
	
	*tree = (nnsliteMTree*)nnsliteMAlloc(sizeof(nnsliteMTree)+(maxnodes-1)*sizeof(nnsliteMTreeNode)+10*1024);
	if (*tree == NULL)
	{
		assert(false);
		return NNSLITE_RET_ERR_ALLOCATE_MEM;
	}
	/* Zero the nodes */
	for (i=0;i<maxnodes; i++)
	{
		NodeZero(&((*tree)->nodes[i]));
	}
	
	(*tree)->objContainerBase = objContainerBase;
	(*tree)->distfunc = func;
	(*tree)->maxDepth = maxdepth;
	(*tree)->maxChildCount = childcount;
	(*tree)->rootNodeIndex = 0;
	(*tree)->maxNumNodes = maxnodes;

	/* Create the root node */
	root = NodeCreate(*tree, true);
	root->depth = 0;
	root->radius = 0.0;
	
	return NNSLITE_RET_OK;
}

NNSLITE_RET nnsliteMTreeFree(nnsliteMTree* tree)
{
	u32 i;
	for (i=0;i<tree->maxNumNodes;i++)
	{
		nnsliteMTreeNode* node = &tree->nodes[i];
		NodeFree(node);
	}
	nnsliteMFree(tree);
	return NNSLITE_RET_OK;
}

NNSLITE_RET nnsliteMTreeInsert(nnsliteMTree* tree, const nnsliteObject* object)
{
	return NodeInsert(tree, tree->rootNodeIndex, object);
}

NNSLITE_RET nnsliteMTreeNNSearch(nnsliteMTree* tree, const nnsliteObject* object, u32 top,nnsliteObject* results)
{
	f64* resultDists = (f64*)nnsliteMAlloc(sizeof(f64)*top);
	u32 i;
	s32 ret;
	
	for (i=0; i<top; i++)
		resultDists[i] = (f64)F64_MAX_VAL;
	
	ret = NodeNNSearch(tree, tree->rootNodeIndex, object, top, results, resultDists);
	NNSLITE_ON_FAILED_RETURN(ret);

	nnsliteMFree(resultDists);
	return NNSLITE_RET_OK;
}


NNSLITE_RET nnsliteMTreeWrite(nnsliteMTree* tree, nnsliteBufferedHeapFile* bhf)
{
	u32 i;
	nnsliteBufferedHeapFileWrite(bhf, tree, sizeof(nnsliteMTree));
	for (i=0;i<tree->maxNumNodes;i++)
	{
		nnsliteMTreeNode* node = &tree->nodes[i];
		nnsliteBufferedHeapFileWrite(bhf, &node->isUsed, sizeof(bool));
		if (node->isUsed)
		{
			NodeWrite(node, bhf);
		}
	}
	return NNSLITE_RET_OK;
}

NNSLITE_RET nnsliteMTreeRead(nnsliteMTree** tree, nnsliteBufferedHeapFile* bhf)
{
	u32 i;
	u32 maxNumNodes;
	nnsliteMTree tempTree;

	nnsliteBufferedHeapFileRead(bhf, &tempTree, sizeof(nnsliteMTree));
	maxNumNodes = tempTree.maxNumNodes;
	/* Allocate the memory for the M-Tree */
	*tree = (nnsliteMTree*)nnsliteMAlloc(sizeof(nnsliteMTree)+(maxNumNodes-1)*sizeof(nnsliteMTreeNode));
	*(*tree) = tempTree;
	/* Read the nodes */
	for(i=0;i<maxNumNodes;i++)
	{
		nnsliteMTreeNode* node = &(*tree)->nodes[i];
		nnsliteBufferedHeapFileRead(bhf, &node->isUsed, sizeof(bool));
		if (node->isUsed)
		{
			NodeRead(node, bhf);
			node->pTree = *tree;
		}
		else
		{
			NodeZero(node);
		}
	}
	return NNSLITE_RET_OK;
}

static NNSLITE_RET NodePrint(nnsliteMTreeNode* node, FILE* fp, u32 indent)
{
	nnsliteMTree* tree = node->pTree;
	nnsliteMTreeNode* childNode;
	u32 childindex;
	u32 i;

	for (i=0;i<indent;i++)
		fprintf(fp,"\t");
	fprintf(fp,"Node ObjectID=%lld, ",node->object.id);
	fprintf(fp,"Radius=%lf, ",node->radius);
	fprintf(fp,"ChildCount=%d, ", node->childCount);

	if (node->isLeaf)
	{
		nnsliteObject obj;
		nnsliteObjectContainerBase* containerBase = tree->objContainerBase;
		nnsliteObjectContainer* container = nnsliteObjectContainerGet(containerBase, node->objContainerID);
		
		
		nnsliteObjectInit(&obj);
		nnsliteObjectContainerBeginScan(container);
		fprintf(fp, "Container (Count=%d): ",container->objCount);
		for (i=0;i<container->objCount;i++)
		{
			nnsliteObjectContainerNext(container, &obj);
			fprintf(fp,"%lld",obj.id);
			if (i < container->objCount-1)
				fprintf(fp,", ");
		}
		fprintf(fp,"\n");
		nnsliteObjectFree(&obj);
	}
	else
	{
		fprintf(fp,"\n");
		for (i=0; i<node->childCount; i++)
		{
			childindex = node->childIndexes[i];
			childNode = &tree->nodes[childindex];
			NodePrint(childNode, fp, indent+1);
		}
	}
	return NNSLITE_RET_OK;
}

NNSLITE_RET nnsliteMTreePrint(nnsliteMTree* tree, FILE* fp)
{
	nnsliteMTreeNode* rootNode = &tree->nodes[tree->rootNodeIndex];
	return NodePrint(rootNode, fp, 0);
}

